#!/usr/bin/python
# -*- coding: UTF-8 -*-
from collections import OrderedDict
from xml.dom.minidom import parse
import xml.dom.minidom
import os
import numpy as np

##
#   Process flow
#       xmls -> data -> txt
##
def xml_parser(path):
    anno = OrderedDict()
    # 使用minidom解析器打开 XML 文档
    DOMTree = xml.dom.minidom.parse(path)
    collection = DOMTree.documentElement

    # 在集合中获取存储目录
    folders = collection.getElementsByTagName("folder")
    # 在集合中获取文件名
    filenames = collection.getElementsByTagName("filename")
    # 获取来源
    sources = collection.getElementsByTagName("source")
    # 获取图片的大小size
    sizes = collection.getElementsByTagName("size")
    # 获取目标object
    objects = collection.getElementsByTagName("object")

    # 打印文件夹的详细信息
    for folder in folders:
        anno['folder'] = folder.firstChild.data

    # 打印文件名的详细信息
    for filename in filenames:
        anno['filename'] = filename.firstChild.data


    # 打印文件来源的详细信息
    for source in sources:
        for s in source.getElementsByTagName("database"):
            anno['database'] =  s.firstChild.data


    # 打印图片的大小的详细信息
    for size in sizes:
        for s in size.getElementsByTagName('width'):
            anno['width'] = s.firstChild.data
        for s in size.getElementsByTagName('height'):
            anno['height'] = s.firstChild.data

    # 打印目标的详细信息
    for obj in objects:
        for s in obj.getElementsByTagName('trackid'):
            anno['trackid'] = s.firstChild.data
        for s in obj.getElementsByTagName('name'):
            anno['name'] = s.firstChild.data
        for s in obj.getElementsByTagName('bndbox'):
            for s_ in s.getElementsByTagName('xmax'):
                anno['xmax'] = s_.firstChild.data
            for s_ in s.getElementsByTagName('xmin'):
                anno['xmin'] = s_.firstChild.data
            for s_ in s.getElementsByTagName('ymax'):
                anno['ymax'] = s_.firstChild.data
            for s_ in s.getElementsByTagName('ymin'):
                anno['ymin'] = s_.firstChild.data
        for s in obj.getElementsByTagName('occluded'):
            anno['occluded'] = s.firstChild.data
        for s in obj.getElementsByTagName('generated'):
            anno['generated'] = s.firstChild.data
    return anno

def read_train_xmls(path):
    data = []
    root_listing_ann = os.listdir(path)
    for ann_name in root_listing_ann:
        full_ann_name = os.path.join(path, ann_name)
        folder_list = os.listdir(full_ann_name)
        folder_list = [os.path.join(full_ann_name, i) for i in folder_list]
        for p in folder_list:
            xmls = os.listdir(p)
            xmls = [os.path.join(p, i) for i in xmls]
            for x in xmls:
                data.append(xml_parser(path=x))

    return data

def read_val_xmls(path):
    data = []
    root_listing_ann = os.listdir(path)
    for ann_name in root_listing_ann:
        full_ann_name = os.path.join(path, ann_name)
        xmls = os.listdir(full_ann_name)
        xmls = [os.path.join(full_ann_name, i) for i in xmls]
        for x in xmls:
            data.append(xml_parser(path=x))

    return data

# 实验发现直接保存成numpy数组太慢而且读取的时候更麻烦
def data2txt(data, prefix='/media/sdc2/Datasets/ILSVRC2015_VID/ILSVRC2015/Data/VID/train'):
    # anno_images = []
    fp = open('./anno_images.txt', 'w+')
    for i in data:
        if 'xmin' in i:
            fp.write(os.path.join(prefix, i['folder'], i['filename'] + '.JPEG '
                                            + i['xmin']+' ' + i['ymin'] + ' '
                                            + i['xmax'] + ' ' + i['ymax']) + '\n')
            #anno_images.append(os.path.join(prefix, i['folder'], i['filename'] + '.JPEG '
            #                                + i['xmin']+' ' + i['ymin'] + ' '
            #                                + i['xmax'] + ' ' + i['ymax']))
    # np.save('anno_images', anno_images)
    fp.close()

def IoU(box, boxes):
    """Compute IoU between detect box and gt boxes

    Parameters:
    ----------
    box: numpy array , shape (5, ): x1, y1, x2, y2, score
        input box
    boxes: numpy array, shape (n, 4): x1, y1, x2, y2
        input ground truth boxes

    Returns:
    -------
    ovr: numpy.array, shape (n, )
        IoU
    """
    box_area = (box[2] - box[0] + 1) * (box[3] - box[1] + 1)
    area = (boxes[2] - boxes[0] + 1) * (boxes[3] - boxes[1] + 1)
    xx1 = np.maximum(box[0], boxes[0])
    yy1 = np.maximum(box[1], boxes[1])
    xx2 = np.minimum(box[2], boxes[2])
    yy2 = np.minimum(box[3], boxes[3])

    # compute the width and height of the bounding box
    w = np.maximum(0, xx2 - xx1 + 1)
    h = np.maximum(0, yy2 - yy1 + 1)

    inter = w * h
    ovr = np.true_divide(inter,(box_area + area - inter))
    #ovr = inter / (box_area + area - inter)
    return ovr

## Test Code
if __name__ == '__main__':
    # path = '/media/sdc2/Datasets/ILSVRC2015_VID/ILSVRC2015/Annotations/VID/'
    prefix = 'D:/Datasets/ILSVRC2015_VID/ILSVRC2015/Data/VID/train'
    # val_data = read_val_xmls(os.path.join(path, 'val'))
    # np.save('val_anno', val_data)
    # train_data = read_train_xmls(os.path.join(path, 'train'))
    # np.save('train_anno', train_data)
    # train_anno = np.load('./anno_data/train_anno.npy')
    # val_anno = np.load('./anno_data/val_anno.npy')
    data = np.load('D:/Datasets/ILSVRC2015_VID/train_anno.npy')
    data2txt(data, prefix=prefix)
    # data = np.load('./train_anno.npy')
    #data2txt(data)

    pass